from skimage import io
import matplotlib.pyplot as plt
import numpy as np
from skimage.restoration import denoise_wavelet
from skimage.metrics import peak_signal_noise_ratio as psnr

clean_img=io.imread("cat.png",as_gray=True)
noisy_img=io.imread("cat_noisy.jpg",as_gray=True)

psnr_list=[]
sigma_value=np.arange(0.01,1,0.01)

for sigma in sigma_value:
    wavelet_img=denoise_wavelet(noisy_img,sigma)
    wavelet_psnr=psnr(clean_img,wavelet_img)
    psnr_list.append(wavelet_psnr)
psnr=np.array(psnr_list)

wavelet_img=denoise_wavelet(noisy_img,sigma_value[np.argmax(psnr)])

fg,ax=plt.subplots(1,4,figsize=(10,4))
ax[0].imshow(clean_img,cmap='gray')
ax[0].set_title("Noise_Free Image",size=8)
ax[1].imshow(noisy_img,cmap='gray')
ax[1].set_title("Noise Image",size=8)
ax[2].imshow(wavelet_img,cmap='gray')
ax[2].set_title("Wavelet Denoised(Best Sigma:{0:0.2f},PSNR:{1:0.2f})".format(sigma_value[np.argmax(psnr)],np.max(psnr)),size=8)
ax[3].plot(sigma_value,psnr)
ax[3].set_title("PSNR vs Sigma",size=8)
corrd="({0:0.2f},{1:0.2f}))".format(sigma_value[np.argmax(psnr)],np.max(psnr))
ax[3].text(sigma_value[np.argmax(psnr)],np.max(psnr),s=corrd)
plt.tight_layout()
plt.show()
